from physics.vector.array import Vector4Array as V4Array
from physics.const import KCHARGED_MASS_GEV as KMASS
from physics.const import PICHARGED_MASS_GEV as PIMASS

import numpy

from charmicat import ColumnCalculator

from cmix.parameters import paramdef

def four_vectors(r):
    v = {}
    v['wsk'] = V4Array(r['wsk_px'],
                       r['wsk_py'],
                       r['wsk_pz'],
                       m=KMASS)
    v['rsk'] = V4Array(r['rsk_px'],
                       r['rsk_py'],
                       r['rsk_pz'],
                       m=KMASS)

    v['rspi'] = V4Array(r['rspi_px'],
                        r['rspi_py'],
                        r['rspi_pz'],
                        m=PIMASS)
    v['wspi'] = V4Array(r['wspi_px'],
                        r['wspi_py'],
                        r['wspi_pz'],
                        m=PIMASS)
    
    v['slow'] = V4Array(r['slow_px'],
                        r['slow_py'],
                        r['slow_pz'],
                        m=PIMASS)
    
    return v    

            

def p_wsk(r, v):
    return v['wsk'].v.mag()

def p_rsk(r, v):
    return v['rsk'].v.mag()

def p_rspi(r, v):
    return v['rspi'].v.mag()

def p_wspi(r, v):
    return v['wspi'].v.mag()

def p_slow(r, v):
    return v['slow'].v.mag()


def rs_dvf_m_pull(r, v):
    num = (r['m_rs'] - paramdef['D0_MASS'])
    denom = numpy.sqrt(numpy.fabs(r['rs_dvf_m_unc']))
    return num/denom

def ws_dvf_m_pull(r, v):
    num = (r['m_ws'] - paramdef['D0_MASS'])
    denom = numpy.sqrt(numpy.fabs(r['ws_dvf_m_unc']))
    return  num/denom


column_calc = ColumnCalculator(
    four_vectors, ("wsk_px wsk_py wsk_pz rsk_px rsk_py rsk_pz " 
                   + "wspi_px wspi_py wspi_pz rspi_px rspi_py rspi_pz "
                   + "slow_px slow_py slow_pz").split())

column_calc.add('p_wsk', p_wsk)
column_calc.add('p_rspi', p_rspi)
column_calc.add('p_rsk', p_rsk)
column_calc.add('p_wspi', p_wspi)
column_calc.add('p_slow', p_slow)
column_calc.add('rs_dvf_m_pull', rs_dvf_m_pull)
column_calc.add('ws_dvf_m_pull', ws_dvf_m_pull)
